// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/dwio/common/unit_loader_tools.h>

using namespace ::testing;
using namespace ::kumo::pollux::dwio::common;
using namespace ::kumo::pollux::dwio::common::unit_loader_tools;

TEST(UnitLoaderToolsTests, NoCallbacksCreated) {
  std::atomic_size_t callCount = 0;
  {
    CallbackOnLastSignal callback([&callCount]() { ++callCount; });
    EXPECT_EQ(callCount, 0);
  }
  EXPECT_EQ(callCount, 1);
}

TEST(UnitLoaderToolsTests, SupportsNullCallbacks) {
  CallbackOnLastSignal callback(nullptr);
  auto cb = callback.getCallback();
  EXPECT_TRUE(cb == nullptr);
}

TEST(UnitLoaderToolsTests, NoExplicitCalls) {
  std::atomic_size_t callCount = 0;
  {
    CallbackOnLastSignal callback([&callCount]() { ++callCount; });
    EXPECT_EQ(callCount, 0);
    {
      auto c1 = callback.getCallback();
      auto c4 = callback.getCallback();
      EXPECT_EQ(callCount, 0);

      auto c2 = std::move(c1);
      auto c3(c2);
      EXPECT_EQ(callCount, 0);

      auto c5 = std::move(c4);
      auto c6(c5);
      EXPECT_EQ(callCount, 0);
    }
    EXPECT_EQ(callCount, 1);
  }
  EXPECT_EQ(callCount, 1);
}

TEST(UnitLoaderToolsTests, NoExplicitCallsFactoryDeletedFirst) {
  std::atomic_size_t callCount = 0;
  {
    std::function<void()> c1, c2;
    {
      CallbackOnLastSignal callback([&callCount]() { ++callCount; });
      EXPECT_EQ(callCount, 0);

      c1 = callback.getCallback();
      c2 = callback.getCallback();
      EXPECT_EQ(callCount, 0);
    }
    EXPECT_EQ(callCount, 0);
  }
  EXPECT_EQ(callCount, 1);
}

TEST(UnitLoaderToolsTests, ExplicitCalls) {
  std::atomic_size_t callCount = 0;
  {
    CallbackOnLastSignal callback([&callCount]() { ++callCount; });
    EXPECT_EQ(callCount, 0);
    {
      auto c1 = callback.getCallback();
      auto c4 = callback.getCallback();
      EXPECT_EQ(callCount, 0);

      c1();
      auto c2 = std::move(c1);
      c2();
      auto c3(c2);
      c3();
      EXPECT_EQ(callCount, 0);

      c4();
      EXPECT_EQ(callCount, 1);
      auto c5 = std::move(c4);
      c5();
      auto c6(c2);
      c6();
      EXPECT_EQ(callCount, 1);
    }
    EXPECT_EQ(callCount, 1);
  }
  EXPECT_EQ(callCount, 1);
}

TEST(UnitLoaderToolsTests, WillOnlyCallbackOnce) {
  std::atomic_size_t callCount = 0;
  {
    CallbackOnLastSignal callback([&callCount]() { ++callCount; });
    EXPECT_EQ(callCount, 0);
    {
      auto c1 = callback.getCallback();
      auto c4 = callback.getCallback();
      EXPECT_EQ(callCount, 0);

      c1();
      auto c2 = std::move(c1);
      c2();
      auto c3(c2);
      c3();
      EXPECT_EQ(callCount, 0);

      c4();
      EXPECT_EQ(callCount, 1);
      auto c5 = std::move(c4);
      c5();
      auto c6(c2);
      c6();
      EXPECT_EQ(callCount, 1);

      // This won't emit a new call
      auto c7 = callback.getCallback();
      c7();
      EXPECT_EQ(callCount, 1);
    }
    EXPECT_EQ(callCount, 1);
  }
  EXPECT_EQ(callCount, 1);
}

TEST(UnitLoaderToolsTests, HowMuchToSkip) {
  // Helpers
  auto testSkip = [](uint64_t rowsToSkip, std::vector<uint64_t> rowCount) {
    return howMuchToSkip(rowsToSkip, rowCount.cbegin(), rowCount.cend());
  };

  auto result = [](uint32_t unitsToSkip, uint64_t rowsToSkip) {
    return std::make_pair(unitsToSkip, rowsToSkip);
  };

  static const char* kErrorMessage =
      "Can't skip more rows than all the rows in all the units";

  // Test cases
  EXPECT_EQ(testSkip(0, {}), result(0, 0));
  POLLUX_ASSERT_THROW(testSkip(1, {}), kErrorMessage);

  EXPECT_EQ(testSkip(0, {0}), result(1, 0));
  POLLUX_ASSERT_THROW(testSkip(1, {0}), kErrorMessage);

  EXPECT_EQ(testSkip(0, {1}), result(0, 0));
  EXPECT_EQ(testSkip(1, {1}), result(1, 0));
  POLLUX_ASSERT_THROW(testSkip(2, {1}), kErrorMessage);

  std::vector<uint64_t> rowCount = {2, 1, 2};
  EXPECT_EQ(testSkip(0, rowCount), result(0, 0));
  EXPECT_EQ(testSkip(1, rowCount), result(0, 1));
  EXPECT_EQ(testSkip(2, rowCount), result(1, 0));
  EXPECT_EQ(testSkip(3, rowCount), result(2, 0));
  EXPECT_EQ(testSkip(4, rowCount), result(2, 1));
  EXPECT_EQ(testSkip(5, rowCount), result(3, 0));
  POLLUX_ASSERT_THROW(testSkip(6, rowCount), kErrorMessage);
}
